import numpy as np

from utils import common
from utils.log import Logger
from src import train
from sklearn.metrics import precision_score, recall_score, roc_auc_score, accuracy_score
import joblib
import os
import pandas as pd
from sklearn.metrics import (
    confusion_matrix, roc_curve, auc, precision_recall_curve,
    average_precision_score, classification_report
)
import matplotlib.pyplot as plt
import pandas as pd
 # 强制优先使用黑体
plt.rcParams['font.family'] = 'SimHei'
plt.rcParams['font.size'] = 15

# 创建日志记录器
logger = Logger(root_path=os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 
                log_name='predict', level='info').get_logger()


def model_predict(path, target_col):
    """
    模型预测
    输入：1.数据集 2.标签列名称
    将输入的数据集进行模型预测
    并输出预测结果
    """
    try:
        logger.info(f"开始模型预测，数据文件路径: {path}，目标列: {target_col}")
        
        # 加载模型
        model_path = '../models/NB_model.pkl'
        model = joblib.load(model_path)
        logger.info(f"成功加载模型: {model_path}")
        
        # 数据预处理
        logger.info("开始数据预处理")
        data, scaler = common.data_preprocessing(path)
        logger.info(f"数据预处理完成，预测数据形状: {data.shape}")
        
        # 分离特征和标签
        y = data[target_col]
        x = data.drop(columns=target_col)
        logger.info(f"特征数量: {x.shape[1]}，预测样本数量: {x.shape[0]}")
        logger.info(f"真实标签分布: {y.value_counts().to_dict()}")
        
        # 进行预测
        logger.info("开始进行预测")
        y_pred = model.predict(x)
        y1_proba = model.predict_proba(x)[:, 1]
        logger.info("预测完成")
        
        # 计算性能指标
        accuracy = accuracy_score(y, y_pred)
        precision = precision_score(y, y_pred)
        recall = recall_score(y, y_pred)
        roc_auc = roc_auc_score(y, y1_proba)
        
        logger.info(f"模型性能指标:")
        logger.info(f"准确率 (Accuracy): {accuracy:.4f}")
        logger.info(f"精确率 (Precision): {precision:.4f}")
        logger.info(f"召回率 (Recall): {recall:.4f}")
        logger.info(f"AUC Score: {roc_auc:.4f}")
        fig, axes = plt.subplots(2, 3, figsize=(18, 12))  # 画布大小：宽18，高12
        fig.suptitle('员工流失分类模型预测结果分析', fontsize=20, fontweight='bold', y=0.98)

        # --------------------------
        # 子图1：混淆矩阵热力图（评估预测正确性）
        # --------------------------
        cm = confusion_matrix(y, y_pred)  # 计算混淆矩阵（TN, FP; FN, TP）
        classes = ['未流失(0)', '已流失(1)']  # 类别名称（根据你的标签含义调整）

        # 绘制热力图
        im = axes[0, 0].imshow(cm, interpolation='nearest', cmap=plt.cm.Blues)
        axes[0, 0].set_title('混淆矩阵', fontsize=14, fontweight='bold', pad=20)
        axes[0, 0].set_xlabel('预测标签', fontsize=12)
        axes[0, 0].set_ylabel('真实标签', fontsize=12)
        axes[0, 0].set_xticks(np.arange(len(classes)))
        axes[0, 0].set_yticks(np.arange(len(classes)))
        axes[0, 0].set_xticklabels(classes)
        axes[0, 0].set_yticklabels(classes)

        # 在热力图上标注具体数值
        for i in range(len(classes)):
            for j in range(len(classes)):
                text = axes[0, 0].text(j, i, cm[i, j], ha='center', va='center',
                                       color='white' if cm[i, j] > cm.max() / 2 else 'black',
                                       fontsize=12, fontweight='bold')

        # 添加颜色条
        cbar = plt.colorbar(im, ax=axes[0, 0], shrink=0.8)
        cbar.set_label('样本数量', fontsize=10)

        # --------------------------
        # 子图2：ROC曲线（评估二分类区分能力）
        # --------------------------
        fpr, tpr, _ = roc_curve(y, y1_proba)  # 计算假正率、真正率
        roc_auc = auc(fpr, tpr)  # 计算AUC值

        axes[0, 1].plot(fpr, tpr, color='darkorange', lw=2,
                        label=f'ROC曲线 (AUC = {roc_auc:.3f})')
        axes[0, 1].plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--',
                        label='随机猜测 (AUC = 0.5)')
        axes[0, 1].set_xlim([0.0, 1.0])
        axes[0, 1].set_ylim([0.0, 1.05])
        axes[0, 1].set_title('ROC曲线', fontsize=14, fontweight='bold', pad=20)
        axes[0, 1].set_xlabel('假正率 (FPR)', fontsize=12)
        axes[0, 1].set_ylabel('真正率 (TPR)', fontsize=12)
        axes[0, 1].legend(loc='lower right', fontsize=10)
        axes[0, 1].grid(True, alpha=0.3)

        # --------------------------
        # 子图3：PR曲线（评估正类预测质量，适合不平衡数据）
        # --------------------------
        precision, recall, _ = precision_recall_curve(y, y1_proba)  # 计算精确率、召回率
        ap = average_precision_score(y, y1_proba)  # 计算平均精度（AP）

        axes[0, 2].plot(recall, precision, color='darkred', lw=2,
                        label=f'PR曲线 (AP = {ap:.3f})')
        # 添加“随机猜测”基准线（正类占比 = 正样本数/总样本数）
        pos_ratio = np.sum(y) / len(y)
        axes[0, 2].plot([0, 1], [pos_ratio, pos_ratio], color='gray', lw=2,
                        linestyle='--', label=f'随机猜测 (AP = {pos_ratio:.3f})')
        axes[0, 2].set_xlim([0.0, 1.0])
        axes[0, 2].set_ylim([0.0, 1.05])
        axes[0, 2].set_title('PR曲线', fontsize=14, fontweight='bold', pad=20)
        axes[0, 2].set_xlabel('召回率 (Recall)', fontsize=12)
        axes[0, 2].set_ylabel('精确率 (Precision)', fontsize=12)
        axes[0, 2].legend(loc='lower left', fontsize=10)
        axes[0, 2].grid(True, alpha=0.3)

        # --------------------------
        # 子图4：真实标签vs预测标签 类别分布对比
        # --------------------------
        # 统计真实/预测标签的类别数量
        y_true_count = pd.Series(y).value_counts().sort_index()  # 真实标签计数
        y_pred_count = pd.Series(y_pred).value_counts().sort_index()  # 预测标签计数
        x_pos = np.arange(len(classes))  # x轴位置
        width = 0.35  # 条形图宽度

        # 绘制分组条形图
        bars1 = axes[1, 0].bar(x_pos - width / 2, y_true_count.values, width,
                               label='真实标签', color='#1f77b4', alpha=0.8)
        bars2 = axes[1, 0].bar(x_pos + width / 2, y_pred_count.values, width,
                               label='预测标签', color='#ff7f0e', alpha=0.8)

        # 添加数值标注
        for bar in bars1:
            height = bar.get_height()
            axes[1, 0].text(bar.get_x() + bar.get_width() / 2., height + 5,
                            f'{int(height)}', ha='center', va='bottom', fontsize=10)
        for bar in bars2:
            height = bar.get_height()
            axes[1, 0].text(bar.get_x() + bar.get_width() / 2., height + 5,
                            f'{int(height)}', ha='center', va='bottom', fontsize=10)

        axes[1, 0].set_title('真实标签与预测标签分布对比', fontsize=14, fontweight='bold', pad=20)
        axes[1, 0].set_xlabel('类别', fontsize=12)
        axes[1, 0].set_ylabel('样本数量', fontsize=12)
        axes[1, 0].set_xticks(x_pos)
        axes[1, 0].set_xticklabels(classes)
        axes[1, 0].legend(fontsize=10)
        axes[1, 0].grid(True, alpha=0.3, axis='y')

        # --------------------------
        # 子图5：预测概率分布直方图（区分真实正负类）
        # --------------------------
        # 提取真实正类（流失）和负类（未流失）的预测概率
        proba_pos = y1_proba[y == 1]  # 真实流失样本的预测概率
        proba_neg = y1_proba[y == 0]  # 真实未流失样本的预测概率

        # 绘制直方图
        axes[1, 1].hist(proba_neg, bins=20, alpha=0.6, color='#1f77b4',
                        label='真实未流失(0)', density=True)
        axes[1, 1].hist(proba_pos, bins=20, alpha=0.6, color='#ff7f0e',
                        label='真实已流失(1)', density=True)
        axes[1, 1].axvline(x=0.5, color='red', linestyle='--', lw=2,
                           label='分类阈值(0.5)')  # 默认分类阈值线

        axes[1, 1].set_title('不同真实类别的预测概率分布', fontsize=14, fontweight='bold', pad=20)
        axes[1, 1].set_xlabel('预测为“已流失”的概率', fontsize=12)
        axes[1, 1].set_ylabel('概率密度', fontsize=12)
        axes[1, 1].legend(fontsize=10)
        axes[1, 1].grid(True, alpha=0.3)
        axes[1, 1].set_xlim([0.0, 1.0])

        # --------------------------
        # 子图6：分类报告条形图（精确率/召回率/F1分数）
        # --------------------------
        # 生成分类报告（转为字典便于提取）
        report = classification_report(y, y_pred, target_names=classes, output_dict=True)
        # 提取0类（未流失）和1类（已流失）的指标
        metrics = ['precision', 'recall', 'f1-score']
        class_0_scores = [report[classes[0]][m] for m in metrics]
        class_1_scores = [report[classes[1]][m] for m in metrics]

        x_pos = np.arange(len(metrics))
        width = 0.35

        # 绘制分组条形图
        bars1 = axes[1, 2].bar(x_pos - width / 2, class_0_scores, width,
                               label=classes[0], color='#1f77b4', alpha=0.8)
        bars2 = axes[1, 2].bar(x_pos + width / 2, class_1_scores, width,
                               label=classes[1], color='#ff7f0e', alpha=0.8)

        # 添加数值标注（保留3位小数）
        for bar in bars1:
            height = bar.get_height()
            axes[1, 2].text(bar.get_x() + bar.get_width() / 2., height + 0.01,
                            f'{height:.3f}', ha='center', va='bottom', fontsize=10)
        for bar in bars2:
            height = bar.get_height()
            axes[1, 2].text(bar.get_x() + bar.get_width() / 2., height + 0.01,
                            f'{height:.3f}', ha='center', va='bottom', fontsize=10)

        axes[1, 2].set_title('各类别性能指标对比', fontsize=14, fontweight='bold', pad=20)
        axes[1, 2].set_xlabel('评估指标', fontsize=12)
        axes[1, 2].set_ylabel('指标分数', fontsize=12)
        axes[1, 2].set_xticks(x_pos)
        axes[1, 2].set_xticklabels(metrics)
        axes[1, 2].set_ylim([0.0, 1.1])
        axes[1, 2].legend(fontsize=10)
        axes[1, 2].grid(True, alpha=0.3, axis='y')

        # --------------------------
        # 调整布局并显示/保存图像
        # --------------------------
        plt.tight_layout(rect=[0, 0, 1, 0.95])  # 预留顶部标题空间
        # 保存图像（可选，分辨率300dpi，适合汇报）
        # plt.savefig('../results/attrition_model_analysis.png', dpi=300, bbox_inches='tight')
        plt.savefig('../data/fig/结果图像2.png')
        plt.show()
        # 预测结果分布
        pred_distribution = pd.Series(y_pred).value_counts().to_dict()
        logger.info(f"预测结果分布: {pred_distribution}")
        
        # 输出详细报告
        report = classification_report(y, y_pred)
        logger.info(f"分类报告:\n{report}")
        
        logger.info("模型预测完成")
        
    except FileNotFoundError as e:
        logger.error(f"模型文件或数据文件未找到: {e}")
        raise
    except KeyError as e:
        logger.error(f"数据中缺少目标列: {e}")
        raise
    except Exception as e:
        logger.error(f"模型预测过程中出现未知错误: {e}")
        raise


if __name__ == '__main__':
    # 测试模型预测
    data, scaler = common.data_preprocessing('../data/train.csv')
    train.model_setup(data, 'Attrition')
    model_predict('../data/test2.csv', 'Attrition')
